
from ffmpy import FFmpeg
from pathlib import Path
import subprocess
import logging
logger = logging.getLogger(__name__)

# 支持的音視頻格式（全局變量）
MEDIA_EXTS = ['.mp4', '.mkv', '.avi', '.mov', '.flv', '.wmv', '.m4v', '.mp3', '.aac', '.wav', '.flac', '.ogg', '.m4a']

def get_video_duration(source_file):
    """ 获取视频的时长（秒） """
    try:
        cmd = [
            'ffprobe', 
            '-v', 'error', 
            '-show_entries', 'format=duration', 
            '-of', 'default=noprint_wrappers=1:nokey=1',
            str(source_file)
        ]
        duration = subprocess.check_output(cmd).strip()
        return float(duration)  # 返回时长（秒）
    except Exception as e:
        logger.error(f"获取视频时长出错: {e}")
        return None

def videosplit(starttime, endtime, source_file, target_folder):
    try:
        source_path = Path(source_file)

        # 如果是文件，直接处理
        if source_path.is_file():
            files_to_process = [source_path]
        # 如果是目录，批量处理所有音視頻文件
        elif source_path.is_dir():
            files_to_process = [f for f in source_path.rglob('*') if f.is_file() and f.suffix.lower() in MEDIA_EXTS]
            if not files_to_process:
                logger.error(f"目录 {source_path} 下没有可处理的媒体文件")
                return
        else:
            logger.error(f"输入路径 {source_path} 不存在")
            return

        for file_path in files_to_process:
            # 如果 target_folder 为空，则使用源文件的同一目录
            if not target_folder:
                out_folder = file_path.parent
            else:
                out_folder = Path(target_folder)
            out_folder.mkdir(parents=True, exist_ok=True)

            # 如果 endtime 为空，获取视频长度并使用作为 endtime
            curr_endtime = endtime
            if curr_endtime is None or curr_endtime == '':
                duration = get_video_duration(file_path)
                if duration is not None:
                    curr_endtime = str(int(duration))
                    print(curr_endtime)

            # 使用源文件的后缀来创建目标文件名
            target_file = f"{out_folder}/{file_path.stem}_st{file_path.suffix}"

            ff = FFmpeg(
                inputs={str(file_path): f'-ss {starttime} -to {curr_endtime}'},
                outputs={str(target_file): '-vcodec copy -acodec copy -y'}
            )
            ff.run()
            logger.info(f"--- 成功分割: {file_path} ---")
    except Exception as e:
        logger.error(f"--- FFmpeg出错: {e} ---")

if __name__ == '__main__':
    starttime = '01:42:00'  # 设置分割起始时间
    endtime = '03:05:00'         # 将 endtime 设置为 None
    source_file = r"F:\Documents\Downloads\Teachers (2009) 09年经典.mp4"
    target_folder = r""      # 指定输出文件夹为空
    videosplit(starttime, endtime, source_file, target_folder)